{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use agents.json with Google Sheets API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install openai google-auth-oauthlib google-auth-httplib2 google-auth"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get Test Auth Credentials\n",
    "1. Go to https://developers.google.com/oauthplayground/\n",
    "2. Find \"Sheets v4\" API\n",
    "3. Select all scopes\n",
    "4. Click on \"Exchange authorization code for tokens\"\n",
    "5. Add values below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "auth_data = {\n",
    "  \"access_token\": \"\",\n",
    "  \"scope\": \"https://www.googleapis.com/auth/spreadsheets\",\n",
    "  \"token_type\": \"Bearer\",\n",
    "  \"expires_in\": 3600,\n",
    "  \"refresh_token\": \"\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from typing import Optional\n",
    "from google.oauth2.credentials import Credentials\n",
    "\n",
    "def load_auth_from_json(auth_json: str, auth_path: Optional[str] = None) -> Credentials:\n",
    "    \"\"\"\n",
    "    Load Google OAuth2 credentials directly from a JSON file.\n",
    "    \n",
    "    Args:\n",
    "        auth_path: Path to the auth JSON file containing tokens\n",
    "        \n",
    "    Returns:\n",
    "        Google OAuth2 Credentials object\n",
    "    \"\"\"\n",
    "    \n",
    "    if auth_path:\n",
    "        with open(auth_path, 'r') as f:\n",
    "            token_data = json.load(f)\n",
    "    else:\n",
    "        token_data = json.loads(auth_json)\n",
    "    \n",
    "    return Credentials(\n",
    "        token=token_data['access_token'],\n",
    "        refresh_token=token_data['refresh_token'],\n",
    "        token_uri='https://oauth2.googleapis.com/token',\n",
    "        client_id='',  # These are not needed for sandbox\n",
    "        client_secret='',\n",
    "        scopes=token_data['scope'].split(' ')\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a Google Sheets Agent\n",
    "This class encapsulates an agent with the following functions.\n",
    "1. Setup Google Auth - This function sets up the Google Sheets API credentials.\n",
    "2. Load Agents.json - This function loads the agents.json file.\n",
    "3. Execute Query - This function executes a natural language query against Google Sheets.\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml\n",
    "from typing import Dict, Any, List\n",
    "\n",
    "from google.oauth2.credentials import Credentials\n",
    "from google_auth_oauthlib.flow import InstalledAppFlow\n",
    "from google.auth.transport.requests import Request\n",
    "from openai import OpenAI\n",
    "\n",
    "from agentsjson.core.models import Flow\n",
    "from agentsjson.core.models.bundle import Bundle\n",
    "from agentsjson.core.models.auth import AuthType, OAuth2AuthConfig\n",
    "import agentsjson.core as core\n",
    "from agentsjson.core.executor import execute_flows\n",
    "from agentsjson.core import ToolFormat\n",
    "\n",
    "class GoogleSheetsAgent:\n",
    "    def __init__(self, auth_json_path: str, openai_api_key: str):\n",
    "        \"\"\"\n",
    "        Initialize the Google Sheets Agent.\n",
    "        \n",
    "        Args:\n",
    "            auth_json_path: Path to the auth JSON file containing tokens\n",
    "            openai_api_key: OpenAI API key\n",
    "        \"\"\"\n",
    "        self.auth_json_path = auth_json_path\n",
    "        self.openai_api_key = openai_api_key\n",
    "        self.google_creds = None\n",
    "        self.bundle = None\n",
    "        self.flows = None\n",
    "        \n",
    "        # Initialize OpenAI client\n",
    "        self.openai_client = OpenAI(api_key=openai_api_key)\n",
    "\n",
    "    def setup_google_auth(self) -> None:\n",
    "        \"\"\"Set up Google OAuth2 authentication using the auth JSON file.\"\"\"\n",
    "        if not os.path.exists(self.auth_json_path):\n",
    "            raise FileNotFoundError(f\"Auth JSON file not found at {self.auth_json_path}\")\n",
    "        \n",
    "        self.google_creds = load_auth_from_json(self.auth_json_path)\n",
    "        print(\"Successfully loaded Google Sheets API credentials\")\n",
    "\n",
    "    def load_agents_json(self, agents_json_path: str) -> None:\n",
    "        \"\"\"Load the agents.json file.\"\"\"\n",
    "        try:\n",
    "            # First load the agents.json content\n",
    "            with open(agents_json_path, 'r') as f:\n",
    "                agents_json_content = json.load(f)\n",
    "            \n",
    "            # Get the OpenAPI spec path from the sources\n",
    "            openapi_path = os.path.join(os.path.dirname(agents_json_path), 'openapi.yaml')\n",
    "            \n",
    "            # Load the OpenAPI spec\n",
    "            with open(openapi_path, 'r') as f:\n",
    "                openapi_content = yaml.safe_load(f)\n",
    "            \n",
    "            # Create the bundle structure\n",
    "            bundle_data = {\n",
    "                \"agentsJson\": agents_json_content,\n",
    "                \"openapi\": openapi_content,\n",
    "                \"operations\": {}  # This will be populated by the Bundle model\n",
    "            }\n",
    "            \n",
    "            self.bundle = Bundle.model_validate(bundle_data)\n",
    "            self.flows = self.bundle.agentsJson.flows\n",
    "            print(\"Successfully loaded agents.json and OpenAPI spec\")\n",
    "        except FileNotFoundError as e:\n",
    "            raise Exception(f\"File not found: {str(e)}\")\n",
    "        except json.JSONDecodeError as e:\n",
    "            raise Exception(f\"Invalid JSON in agents.json: {str(e)}\")\n",
    "        except Exception as e:\n",
    "            raise Exception(f\"Error loading agents.json: {str(e)}\")\n",
    "\n",
    "    def execute_query(self, query: str, flow_hint: Optional[List[str]] = None) -> Dict[str, Any]:\n",
    "        \"\"\"Execute a natural language query against Google Sheets.\"\"\"\n",
    "        if not self.google_creds or not self.bundle or not self.flows:\n",
    "            raise Exception(\"Please call setup_google_auth() and load_agents_json() first\")\n",
    "\n",
    "        # Format flows for the prompt\n",
    "        flows = self.flows        \n",
    "        if flow_hint:\n",
    "            flows = [flow for flow in self.flows if flow.id in flow_hint]\n",
    "                    \n",
    "        flows_context = core.flows_prompt(flows)\n",
    "        \n",
    "        # Create system prompt\n",
    "        system_prompt = f\"\"\"You are an AI assistant that helps users interact with the Google Sheets API.\n",
    "You have access to the following API flows:\n",
    "\n",
    "{flows_context}\n",
    "\n",
    "Analyze the user's request and use the appropriate API flows to accomplish the task.\n",
    "You must give your arguments for the tool call as Structured Outputs JSON with keys `parameters` and `requestBody`.\n",
    "When creating or updating spreadsheets, make sure to format the data in a clear and organized way.\n",
    "\n",
    "For the spreadsheets_create flow, provide the requestBody in this format:\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "        tools = core.flows_tools(flows, format=ToolFormat.OPENAI)\n",
    "        print(\"TOOLS\", tools)\n",
    "\n",
    "        # Get completion from OpenAI\n",
    "        response = self.openai_client.chat.completions.create(\n",
    "            model=\"gpt-4o\",\n",
    "            messages=[\n",
    "                {\"role\": \"system\", \"content\": system_prompt},\n",
    "                {\"role\": \"user\", \"content\": query}\n",
    "            ],\n",
    "            tools=core.flows_tools(flows, format=ToolFormat.OPENAI),\n",
    "            temperature=0\n",
    "        )\n",
    "        \n",
    "        # Configure auth\n",
    "        auth = OAuth2AuthConfig(\n",
    "            type=AuthType.OAUTH2,\n",
    "            credentials=self.google_creds,\n",
    "            token=self.google_creds.token,\n",
    "            scopes=set(self.google_creds.scopes),\n",
    "            refresh_token=self.google_creds.refresh_token,\n",
    "        )\n",
    "        \n",
    "        # Execute the flows\n",
    "        result = execute_flows(\n",
    "            response,\n",
    "            format=ToolFormat.OPENAI,\n",
    "            bundle=self.bundle,\n",
    "            flows=flows,\n",
    "            auth=auth\n",
    "        )\n",
    "        \n",
    "        return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "current_dir = os.path.dirname(\"/Users/kaushik/Documents/OceanWater/agents-json/examples/sheets.ipynb\")\n",
    "auth_json_path = os.path.join(current_dir, \"sheetsauth.json\")\n",
    "agents_json_path = os.path.join(current_dir, \"..\", \"agents_json\", \"googlesheets\", \"agents.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "openai_api_key = os.getenv(\"OPENAI_API_KEY\")\n",
    "agent = GoogleSheetsAgent(auth_json_path, openai_api_key)\n",
    "agent.setup_google_auth()\n",
    "agent.load_agents_json(agents_json_path)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"Create a new spreadsheet titled 'Customer Relationship Management' and add the columns 'Name', 'Email', 'Phone', 'Address', 'City', 'State', 'Zip', 'Country'. Then add 10 rows of sample data to it.\"\n",
    "# flow_hint = [] - Optional to help the LLM choose the right flow\n",
    "\n",
    "result = agent.execute_query(query)\n",
    "result"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "newenv2",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
